#!/usr/bin/env python3
"""Probe which request attributes are part of the CloudFront cache key.

Technique: carve out a private cache entry with a unique query-string token
(query string is in the cache key on both configs), warm it with value A until
it returns a HIT, then flip to value B. A MISS on B means the attribute
contributes to the cache key; a HIT means it does not.
"""

import http.client
import random
import time

UA = "crates-cache-test (tobias@bieniek.cloud)"
PATH = "/_app/immutable/chunks/kNaey6uv.js"
DOMAINS = ["cloudfront-app.crates.io", "cloudfront-app.staging.crates.io"]


def get_xcache(domain, qs, headers):
    """Return the X-Cache header value for one request."""
    conn = http.client.HTTPSConnection(domain, timeout=20)
    h = {"User-Agent": UA}
    h.update(headers)
    conn.request("GET", f"{PATH}?{qs}", headers=h)
    resp = conn.getresponse()
    resp.read()
    xc = resp.getheader("X-Cache", "<none>")
    conn.close()
    return xc


def classify(b):
    if b.startswith("Miss"):
        return "IN  cache key"
    if b.startswith("Hit") or b.startswith("RefreshHit"):
        return "NOT in cache key"
    return f"?? ({b})"


def probe(name, a_headers, b_headers):
    print(f"== {name} ==")
    for domain in DOMAINS:
        qs = f"probe-{random.randint(0, 10**12)}"
        a1 = get_xcache(domain, qs, a_headers)
        time.sleep(1)
        a2 = get_xcache(domain, qs, a_headers)
        b = get_xcache(domain, qs, b_headers)
        print(f"  {domain:<34} A=[{a1} | {a2}]  B=[{b}]  -> {classify(b)}")
    print()


probe("Accept", {"Accept": "application/json"}, {"Accept": "text/html"})
probe("Referer", {"Referer": "https://a.example/"}, {"Referer": "https://b.example/"})
# both User-Agent values must be valid custom UAs, else crates.io returns 403
probe(
    "User-Agent",
    {"User-Agent": "crates-cache-test-A (tobias@bieniek.cloud)"},
    {"User-Agent": "crates-cache-test-B (tobias@bieniek.cloud)"},
)
probe("X-Request-Id", {"X-Request-Id": "aaaaaaaa"}, {"X-Request-Id": "bbbbbbbb"})
probe(
    "Authorization",
    {"Authorization": "Bearer tokenA"},
    {"Authorization": "Bearer tokenB"},
)
probe("Cookie", {"Cookie": "probe=a"}, {"Cookie": "probe=b"})
probe("Accept-Encoding", {"Accept-Encoding": "gzip"}, {"Accept-Encoding": "br"})